import os
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.applications import DenseNet121
from tensorflow.keras.layers import GlobalAveragePooling2D, Dense,Dropout
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt
import matplotlib.font_manager as font_manager
import time


# 检查 GPU 是否可用po
physical_devices = tf.config.list_physical_devices('GPU')
print("可用的 GPU 设备:", physical_devices)

# 图像生成器配置
train_datagen = ImageDataGenerator(
    rescale=1.0 / 255.0,
)

test_datagen = ImageDataGenerator(rescale=1.0 / 255.0)

# 加载训练集图像
train_generator = train_datagen.flow_from_directory(
    'Final/data/train',
    target_size=(224, 224),
)

# 加载测试集图像
test_generator = test_datagen.flow_from_directory(
    'Final/data/test',
    target_size=(224, 224),
    batch_size=32,
)

# 构建模型  
with tf.device('/GPU:0'):
    base_model = DenseNet121(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
    x = base_model.output
    x = Dropout(0.01)(x)                                                   
    x = GlobalAveragePooling2D()(x)
    x = Dense(1024, activation='relu')(x)
    predictions = Dense(6, activation='softmax')(x)

    model = Model(inputs=base_model.input, outputs=predictions)


# 编译模型
model.compile(optimizer="Adam", loss='categorical_crossentropy', metrics=['accuracy'])

# 模型训练
with tf.device('/GPU:0'):
    history = model.fit(
        train_generator,
        steps_per_epoch=len(train_generator),
        epochs=100,
        validation_data=test_generator
    )

loss_values = history.history['loss']
accuracy_values = history.history['accuracy']
val_loss_values = history.history['val_loss']
val_accuracy_values = history.history['val_accuracy']

plt.plot(loss_values, label='Training Loss')
plt.plot(val_loss_values, label='Validation Loss')
plt.title('Total_Loss',fontsize=20)
plt.xlabel('Epochs', fontsize=16)
plt.ylabel('Loss', fontsize=16)
plt.legend(prop={'size': 14},frameon=False)
plt.show()
plt.clf()
# 绘制训练准确率图像
plt.plot(accuracy_values, label='Training Accuracy')
plt.plot(val_accuracy_values, label='Validation Accuracy')
plt.title('Total_Accuracy',fontsize=16)
plt.xlabel('Epochs', fontsize=16)
plt.ylabel('Accuracy', fontsize=16)
plt.show()

# 保存模型
model.save('DenseNet121.h5')



